using System;
using System.Collections;
using System.Collections.Generic;
using System.IO;
using System.Linq;
using System.Runtime.CompilerServices;
using System.Runtime.InteropServices;
using System.Text;
using System.Threading;
using WindowsAPI.Service;

namespace WindowsAPI
{
    [CopyRight("https://vmccontroller.svn.codeplex.com/svn/VmcController/VmcServices/DetectOpenFiles.cs")]
    [ComVisible(true)]//, EventTrackingEnabled(true)]
    public class DetectOpenFiles// : ServicedComponent
    {
        private static Dictionary<string, string> deviceMap;
        private const string networkDevicePrefix = "\\Device\\LanmanRedirector\\";

        private enum SystemHandleType
        {
            // ReSharper disable UnusedMember.Local
            OB_TYPE_UNKNOWN = 0,
            OB_TYPE_TYPE = 1,
            OB_TYPE_DIRECTORY,
            OB_TYPE_SYMBOLIC_LINK,
            OB_TYPE_TOKEN,
            OB_TYPE_PROCESS,
            OB_TYPE_THREAD,
            OB_TYPE_UNKNOWN_7,
            OB_TYPE_EVENT,
            OB_TYPE_EVENT_PAIR,
            OB_TYPE_MUTANT,
            OB_TYPE_UNKNOWN_11,
            OB_TYPE_SEMAPHORE,
            OB_TYPE_TIMER,
            OB_TYPE_PROFILE,
            OB_TYPE_WINDOW_STATION,
            OB_TYPE_DESKTOP,
            OB_TYPE_SECTION,
            OB_TYPE_KEY,
            OB_TYPE_PORT,
            OB_TYPE_WAITABLE_PORT,
            OB_TYPE_UNKNOWN_21,
            OB_TYPE_UNKNOWN_22,
            OB_TYPE_UNKNOWN_23,
            OB_TYPE_UNKNOWN_24,
            //OB_TYPE_CONTROLLER,
            //OB_TYPE_DEVICE,
            //OB_TYPE_DRIVER,
            OB_TYPE_IO_COMPLETION,
            OB_TYPE_FILE
            // ReSharper restore UnusedMember.Local
        };

        private const int handleTypeTokenCount = 27;
        private static readonly string[] handleTypeTokens = new[] 
        { 
            "", "", "Directory", "SymbolicLink", "Token", "Process", "Thread", "Unknown7", "Event", "EventPair", "Mutant",
            "Unknown11", "Semaphore", "Timer", "Profile", "WindowStation", "Desktop", "Section", "Key", "Port", "WaitablePort",
            "Unknown21", "Unknown22", "Unknown23", "Unknown24", "IoCompletion", "File"
        };

        [StructLayout(LayoutKind.Sequential)]
        private struct SYSTEM_HANDLE_ENTRY
        {
            // ReSharper disable FieldCanBeMadeReadOnly.Local
            public int OwnerPid;
            public byte ObjectType;
            public byte HandleFlags;
            public short HandleValue;
            public int ObjectPointer;
            public int AccessMask;
            // ReSharper restore FieldCanBeMadeReadOnly.Local
        }

        /// <summary>Gets the open files enumerator.</summary>
        /// <param name="processId">The process id.</param>
        /// <returns></returns>
        public static IEnumerator<FileSystemInfo> GetOpenFilesEnumerator(int processId)
        {
            return new OpenFiles(processId).GetEnumerator();
        }

        private sealed class OpenFiles : IEnumerable<FileSystemInfo>
        {
            private readonly int f_PID;

            internal OpenFiles(int PID) { f_PID = PID; }

            #region IEnumerable<FileSystemInfo> Members

            public IEnumerator<FileSystemInfo> GetEnumerator()
            {
                NT_STATUS ret;
                var length = 0x10000;
                // Loop, probing for required memory.

                do
                {
                    var ptr = IntPtr.Zero;
                    RuntimeHelpers.PrepareConstrainedRegions();
                    try
                    {
                        RuntimeHelpers.PrepareConstrainedRegions();
                        try { }
                        finally
                        {
                            // CER guarantees that the address of the allocated 
                            // memory is actually assigned to ptr if an 
                            // asynchronous exception occurs.
                            ptr = Marshal.AllocHGlobal(length);
                        }

                        int returnLength;
                        ret = ntdll.NtQuerySystemInformation(SYSTEM_INFORMATION_CLASS.SystemHandleInformation, ptr, length, out returnLength);

                        switch(ret)
                        {
                            case NT_STATUS.InfoLengthMismatch:
                                // Round required memory up to the nearest 64KB boundary.
                                length = ((returnLength + 0xffff) & ~0xffff);
                                break;
                            case NT_STATUS.Success:
                                var handleCount = Marshal.ReadInt32(ptr);
                                var offset = sizeof(int);
                                var size = Marshal.SizeOf(typeof(SYSTEM_HANDLE_ENTRY));
                                for(var i = 0; i < handleCount; i++)
                                {
                                    var handleEntry =
                                                (SYSTEM_HANDLE_ENTRY)Marshal.PtrToStructure((IntPtr)((int)ptr + offset),
                                                                                            typeof(SYSTEM_HANDLE_ENTRY));
                                    if(handleEntry.OwnerPid == f_PID)
                                    {
                                        var handle = (IntPtr)handleEntry.HandleValue;
                                        SystemHandleType handleType;

                                        if(GetHandleType(handle, handleEntry.OwnerPid, out handleType)
                                           && handleType == SystemHandleType.OB_TYPE_FILE)
                                        {
                                            string devicePath;
                                            if(GetFileNameFromHandle(handle, handleEntry.OwnerPid, out devicePath))
                                            {
                                                string dosPath;
                                                if(ConvertDevicePathToDosPath(devicePath, out dosPath))
                                                {
                                                    if(File.Exists(dosPath))
                                                        yield return new FileInfo(dosPath);
                                                    else if(Directory.Exists(dosPath))
                                                        yield return new DirectoryInfo(dosPath);
                                                }
                                            }
                                        }
                                    }
                                    offset += size;
                                }
                                // ReSharper disable RedundantJumpStatement
                                break;
                            // ReSharper restore RedundantJumpStatement
                        }
                    }
                    finally
                    {
                        // CER guarantees that the allocated memory is freed, 
                        // if an asynchronous exception occurs. 
                        Marshal.FreeHGlobal(ptr);
                        //sw.Flush();
                        //sw.Close();
                    }
                }
                while(ret == NT_STATUS.InfoLengthMismatch);
            }

            #endregion

            #region IEnumerable Members

            IEnumerator IEnumerable.GetEnumerator() { return GetEnumerator(); }

            #endregion
        }

        #region Private Members

        private static bool GetFileNameFromHandle(IntPtr handle, int processId, out string fileName)
        {
            var currentProcess = Kernel32.GetCurrentProcess();
            var remote = (processId != Kernel32.GetProcessId(currentProcess));
            SafeProcessHandle processHandle = null;
            SafeObjectHandle objectHandle = null;
            try
            {
                if(remote)
                {
                    processHandle = Kernel32.OpenProcess(ProcessAccessTypes.PROCESS_DUP_HANDLE, true, processId);
                    if(Kernel32.DuplicateHandle(processHandle.DangerousGetHandle(), handle, currentProcess,
                            out objectHandle, 0, false, DuplicateHandleOptions.DUPLICATE_SAME_ACCESS))
                        handle = objectHandle.DangerousGetHandle();
                }
                return GetFileNameFromHandle(handle, out fileName, 200);
            }
            finally
            {
                if(remote)
                {
                    if(processHandle != null)
                        processHandle.Close();
                    if(objectHandle != null)
                        objectHandle.Close();
                }
            }
        }
        private static bool GetFileNameFromHandle(IntPtr handle, out string fileName, int wait)
        {
            using(var f = new FileNameFromHandleState(handle))
            {
                ThreadPool.QueueUserWorkItem(GetFileNameFromHandle, f);
                if(f.WaitOne(wait))
                {
                    fileName = f.FileName;
                    return f.RetValue;
                }

                fileName = string.Empty;
                return false;
            }
        }

        private class FileNameFromHandleState : IDisposable
        {
            private readonly ManualResetEvent f_ManualResetEvent;

            private readonly IntPtr f_Handle;

            public IntPtr Handle { get { return f_Handle; } }

            public string FileName { get; set; }

            public bool RetValue { get; set; }

            public FileNameFromHandleState(IntPtr handle)
            {
                f_ManualResetEvent = new ManualResetEvent(false);
                f_Handle = handle;
            }

            public bool WaitOne(int wait) { return f_ManualResetEvent.WaitOne(wait, false); }

            public void Set() { f_ManualResetEvent.Set(); }

            #region IDisposable Members

            public void Dispose() { if(f_ManualResetEvent != null)  f_ManualResetEvent.Close(); }

            #endregion
        }

        private static void GetFileNameFromHandle(object state)
        {
            var s = (FileNameFromHandleState)state;
            string fileName;
            s.RetValue = GetFileNameFromHandle(s.Handle, out fileName);
            s.FileName = fileName;
            s.Set();
        }

        private static bool GetFileNameFromHandle(IntPtr handle, out string fileName)
        {
            var ptr = IntPtr.Zero;
            RuntimeHelpers.PrepareConstrainedRegions();
            try
            {
                var length = 512;  // 512 bytes
                RuntimeHelpers.PrepareConstrainedRegions();
                try { }
                finally
                {
                    // CER guarantees the assignment of the allocated 
                    // memory address to ptr, if an ansynchronous exception 
                    // occurs.
                    ptr = Marshal.AllocHGlobal(length);
                }

                var ret = ntdll.NtQueryObject(handle, OBJECT_INFORMATION_CLASS.ObjectNameInformation, ptr, length, out length);
                if(ret == NT_STATUS.Buffer_Overflow)
                {
                    RuntimeHelpers.PrepareConstrainedRegions();
                    try { }
                    finally
                    {
                        // CER guarantees that the previous allocation is freed,
                        // and that the newly allocated memory address is 
                        // assigned to ptr if an asynchronous exception occurs.
                        Marshal.FreeHGlobal(ptr);
                        ptr = Marshal.AllocHGlobal(length);
                    }
                    ret = ntdll.NtQueryObject(handle, OBJECT_INFORMATION_CLASS.ObjectNameInformation, ptr, length, out length);
                }
                if(ret == NT_STATUS.Success)
                {
                    fileName = Marshal.PtrToStringUni((IntPtr)((int)ptr + 8), (length - 9) / 2);
                    return fileName.Length != 0;
                }
            }
            finally
            {
                // CER guarantees that the allocated memory is freed, 
                // if an asynchronous exception occurs.
                Marshal.FreeHGlobal(ptr);
            }

            fileName = string.Empty;
            return false;
        }

        private static bool GetHandleType(IntPtr handle, int processId, out SystemHandleType handleType)
        {
            return GetHandleTypeFromToken(GetHandleTypeToken(handle, processId), out handleType);
        }

        /*
        private static bool GetHandleType(IntPtr handle, out SystemHandleType handleType)
        {
            //var token = GetHandleTypeToken(handle);
            return GetHandleTypeFromToken(GetHandleTypeToken(handle), out handleType);
        }
        */

        private static bool GetHandleTypeFromToken(string token, out SystemHandleType handleType)
        {
            for(var i = 1; i < handleTypeTokenCount; i++)
            {
                if(handleTypeTokens[i] != token) continue;
                handleType = (SystemHandleType)i;
                return true;
            }

            handleType = SystemHandleType.OB_TYPE_UNKNOWN;
            return false;
        }

        private static string GetHandleTypeToken(IntPtr handle, int processId)
        {
            var currentProcess = Kernel32.GetCurrentProcess();
            var remote = (processId != Kernel32.GetProcessId(currentProcess));
            SafeProcessHandle processHandle = null;
            SafeObjectHandle objectHandle = null;
            try
            {
                if(remote)
                {
                    processHandle = Kernel32.OpenProcess(ProcessAccessTypes.PROCESS_DUP_HANDLE, true, processId);
                    if(Kernel32.DuplicateHandle(processHandle.DangerousGetHandle(), handle, currentProcess,
                                                out objectHandle, 0, false, DuplicateHandleOptions.DUPLICATE_SAME_ACCESS))
                        handle = objectHandle.DangerousGetHandle();
                }
                return GetHandleTypeToken(handle);
            }
            finally
            {
                if(remote)
                {
                    if(processHandle != null)
                        processHandle.Close();
                    if(objectHandle != null)
                        objectHandle.Close();
                }
            }
        }

        private static string GetHandleTypeToken(IntPtr handle)
        {
            int length;
            ntdll.NtQueryObject(handle, OBJECT_INFORMATION_CLASS.ObjectTypeInformation, IntPtr.Zero, 0, out length);
            var ptr = IntPtr.Zero;
            RuntimeHelpers.PrepareConstrainedRegions();
            try
            {
                RuntimeHelpers.PrepareConstrainedRegions();
                try { }
                finally { ptr = Marshal.AllocHGlobal(length); }

                if(ntdll.NtQueryObject(handle, OBJECT_INFORMATION_CLASS.ObjectTypeInformation, ptr, length, out length)
                   == NT_STATUS.Success)
                    return Marshal.PtrToStringUni((IntPtr)((int)ptr + 0x60));
            }
            finally { Marshal.FreeHGlobal(ptr); }
            return string.Empty;
        }

        public static bool ConvertDevicePathToDosPath(string devicePath, out string dosPath)
        {
            EnsureDeviceMap();
            var i = devicePath.Length;
            while(i > 0 && (i = devicePath.LastIndexOf('\\', i - 1)) != -1)
            {
                string drive;
                if(!deviceMap.TryGetValue(devicePath.Substring(0, i), out drive)) continue;
                dosPath = string.Concat(drive, devicePath.Substring(i));
                return dosPath.Length != 0;
            }
            dosPath = string.Empty;
            return false;
        }

        private static void EnsureDeviceMap()
        {
            if(deviceMap != null) return;

            var localDeviceMap = BuildDeviceMap();
            Interlocked.CompareExchange(ref deviceMap, localDeviceMap, null);
        }

        private static Dictionary<string, string> BuildDeviceMap()
        {
            var logicalDrives = Environment.GetLogicalDrives();
            var localDeviceMap = new Dictionary<string, string>(logicalDrives.Length);
            var lpTargetPath = new StringBuilder(Consts.MAX_PATH);

            foreach(var lpDeviceName in logicalDrives.Select(lpDevice => lpDevice.Substring(0, 2)))
            {
                Kernel32.QueryDosDevice(lpDeviceName, lpTargetPath, Consts.MAX_PATH);
                localDeviceMap.Add(NormalizeDeviceName(lpTargetPath.ToString()), lpDeviceName);
            }

            localDeviceMap.Add(networkDevicePrefix.Substring(0, networkDevicePrefix.Length - 1), "\\");
            return localDeviceMap;
        }

        private static string NormalizeDeviceName(string deviceName)
        {
            return string.Compare(deviceName, 0, networkDevicePrefix, 0, networkDevicePrefix.Length,
                                  StringComparison.InvariantCulture) == 0
                               ? string.Concat(networkDevicePrefix,
                                               deviceName.Substring(deviceName.IndexOf('\\', networkDevicePrefix.Length) + 1))
                               : deviceName;
        }

        #endregion
    }
}